from threading import Thread
from utils.common import exec_shell
import time

def collect_cpu_utility_or_memory_usage(fp,period,cmd):
    while True:
        output,_=exec_shell(cmd)
        fp.write(output.decode()+',')
        time.sleep(period)

def collect_cpu_utility_and_memory_usage(cu_fp,mu_fp,period,cmd):
    while True:
        output,_=exec_shell(cmd)
        cpu,mem=output.decode().split()
        cu_fp.write(cpu+',')
        mu_fp.write(mem+',')
        time.sleep(period)

def collect_gpu_info(fp,period,cmd):
    while True:
        output,_=exec_shell(cmd)
        fp.write(output.decode()+'\n')
        time.sleep(period)

def collect_hardware_cost_metric_per_period(conn,bttask_json,pid,write_api):
    threads=[]
    samples_dir='btresults/'+bttask_json['name']
    cpu_utility_samples_fp=None
    memory_usage_samples_fp=None
    gpu_info_samples_fp=None
    b_exec=False
    if all(x in bttask_json['metrics'] for x in ('cpu_utility','memory_usage')):
        b_exec=True
        cpu_utility_samples_fp=open(samples_dir+'/cpu_utility_samples.txt','a')
        memory_usage_samples_fp=open(samples_dir+'/memory_usage_samples.txt','a')

        threads.append(Thread(target=collect_cpu_utility_and_memory_usage,args=(cpu_utility_samples_fp,memory_usage_samples_fp,bttask_json['hardware_cost_collection_interval'],f"top -bn 1 -p {pid}|tail -3|tail -1|awk '{{cpu=NF-3}} {{mem=NF-6}} {{print $cpu\" \"$mem}}'")))
    elif any(x in bttask_json['metrics'] for x in ('cpu_utility','memory_usage')):
        b_exec=True
        if 'memory_usage' in bttask_json['metrics']:
            memory_usage_samples_fp=open(samples_dir+'/memory_usage_samples.txt','a')
            threads.append(Thread(target=collect_cpu_utility_or_memory_usage,args=(memory_usage_samples_fp,bttask_json['hardware_cost_collection_interval'],f"top -bn 1 -p {pid}|tail -3|tail -1|awk '{{mem=NF-6}} {{print $mem}}'")))
        else:
            cpu_utility_samples_fp=open(samples_dir+'/cpu_utility_samples.txt','a')
            threads.append(Thread(target=collect_cpu_utility_or_memory_usage,args=(cpu_utility_samples_fp,bttask_json['hardware_cost_collection_interval'],f"top -bn 1 -p {pid}|tail -3|tail -1|awk '{{cpu=NF-3}} {{print $cpu}}'")))

    # https://www.cnblogs.com/wsnan/p/11769838.html 服务主机尽量不要运行显示程序，获取GPU利用率和显存利用率是总体的情况，并没有分出单独的某一进程占用的GPU利用率、显存利用率和显存使用量这种信息，好像还得动底层，我又不会
    # nvidia-smi --query-gpu=utilization.gpu,utilization.memory,memory.used,temperature.gpu,power.draw,clocks.sm --format=csv,noheader
    query_gpu_str=''
    if 'gpu_utility' in bttask_json['metrics']:
        query_gpu_str+='utilization.gpu,'
    if 'gpu_memory_utility' in bttask_json['metrics']:
        query_gpu_str+='utilization.memory,'
    if 'gpu_memory_usage' in bttask_json['metrics']:
        query_gpu_str+='memory.used,'
    if 'gpu_temperature' in bttask_json['metrics']:
        query_gpu_str+='temperature.gpu,'
    if 'gpu_power' in bttask_json['metrics']:
        query_gpu_str+='power.draw,'
    if 'gpu_clock_frequency' in bttask_json['metrics']:
        query_gpu_str+='clocks.sm,'
    if query_gpu_str!='':
        b_exec=True
        query_gpu_str+='timestamp'
        gpu_info_samples_fp=open(samples_dir+'/gpu_info_samples.txt','a')
        threads.append(Thread(target=collect_gpu_info,args=(gpu_info_samples_fp,bttask_json['hardware_cost_collection_interval'],f'nvidia-smi --query-gpu={query_gpu_str} --format=csv,noheader')))

    if b_exec:
        conn.send('executing')
        for t in threads:
            t.daemon=True
            t.start() # 本进程退出，子线程均退出
    else:
        conn.send('idle')

    if cpu_utility_samples_fp is not None:
        cpu_utility_samples_fp.close()
    if memory_usage_samples_fp is not None:
        memory_usage_samples_fp.close()
    if gpu_info_samples_fp is not None:
        gpu_info_samples_fp.close()

if __name__=='__main__':
    pass